#include "VulkanTool.h"
#include "RenderState.h"

TextureTransfer::TextureTransfer()
{
	mImageMemoryBarrier = {};
	mImageMemoryBarrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
	mImageMemoryBarrier.srcAccessMask = 0;
	mImageMemoryBarrier.dstAccessMask = 0;
	mImageMemoryBarrier.oldLayout = VK_IMAGE_LAYOUT_MAX_ENUM;
	mImageMemoryBarrier.newLayout = VK_IMAGE_LAYOUT_MAX_ENUM;
	mImageMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	mImageMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	mImageMemoryBarrier.subresourceRange = { VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM, 0, VK_REMAINING_MIP_LEVELS, 0, VK_REMAINING_ARRAY_LAYERS };

	mSrcStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
	mDstStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
}

VkImageMemoryBarrier TextureTransfer::GetImageMemoryBarrier()
{
	VkImageMemoryBarrier imageMemoryBarrier = mImageMemoryBarrier;
	if (imageMemoryBarrier.oldLayout == VK_IMAGE_LAYOUT_MAX_ENUM)
	{
		imageMemoryBarrier.oldLayout = mTexture->GetImageLayout();
	}
	if (imageMemoryBarrier.subresourceRange.aspectMask == VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM)
	{
		imageMemoryBarrier.subresourceRange = mTexture->GetSubresourceRange();
	}
	imageMemoryBarrier.image = mTexture->GetImage();
	mTexture->mImageLayouts[mTexture->GetFrameIndex()] = imageMemoryBarrier.newLayout;
	return imageMemoryBarrier;
}

TextureTransfer& TextureTransfer::SetTexture(RefCountPtr<RHITexture> texture) noexcept
{
	mTexture = texture;
	return *this;
}

TextureTransfer& TextureTransfer::SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept
{
	mSrcStageMask = srcStageMask;
	return *this;
}

TextureTransfer& TextureTransfer::SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept
{
	mDstStageMask = dstStageMask;
	return *this;
}

TextureTransfer& TextureTransfer::SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept
{
	mImageMemoryBarrier.srcAccessMask = srcAccessMask;
	return *this;
}

TextureTransfer& TextureTransfer::SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept
{
	mImageMemoryBarrier.dstAccessMask = dstAccessMask;
	return *this;
}

TextureTransfer& TextureTransfer::SetOldLayout(VkImageLayout oldLayout) noexcept
{
	mImageMemoryBarrier.oldLayout = oldLayout;
	return *this;
}

TextureTransfer& TextureTransfer::SetNewLayout(VkImageLayout newLayout) noexcept
{
	mImageMemoryBarrier.newLayout = newLayout;
	return *this;
}

TextureTransfer& TextureTransfer::SetSrcQueueFamily(VkQueueFlagBits srcQueueFlag) noexcept
{
	mImageMemoryBarrier.srcQueueFamilyIndex = GVulkanRHI->GetQueueFamily(srcQueueFlag);
	return *this;
}

TextureTransfer& TextureTransfer::SetDstQueueFamily(VkQueueFlagBits dstQueueFlag) noexcept
{
	mImageMemoryBarrier.dstQueueFamilyIndex = GVulkanRHI->GetQueueFamily(dstQueueFlag);
	return *this;
}

TextureTransfer& TextureTransfer::SetSubresourceRange(VkImageSubresourceRange subresourceRange) noexcept
{
	mImageMemoryBarrier.subresourceRange = subresourceRange;
	return *this;
}

void TextureTransfer::Transfer(VkQueueFlagBits queueFlag /*= VK_QUEUE_GRAPHICS_BIT*/, bool isMultiple /*= false*/)
{
	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand(queueFlag);

	Transfer(commandBuffer, isMultiple);

	GRenderState->FlushSingleCommand(commandBuffer);
}

void TextureTransfer::Transfer(VkCommandBuffer commandBuffer, bool isMultiple /*= false*/)
{
	/*assert(mSrcStageMask != VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM);
	assert(mDstStageMask != VK_PIPELINE_STAGE_FLAG_BITS_MAX_ENUM);
	assert(mImageMemoryBarrier.srcAccessMask != VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM);
	assert(mImageMemoryBarrier.dstAccessMask != VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM);*/
	assert(mImageMemoryBarrier.newLayout != VK_IMAGE_LAYOUT_MAX_ENUM);

	if (isMultiple && mTexture->IsMultiple())
	{
		for (uint32 i = 0; i < GRenderState->GetFrameCount(); ++i)
		{
			VkImageMemoryBarrier imageMemoryBarrier = mImageMemoryBarrier;
			if (imageMemoryBarrier.oldLayout == VK_IMAGE_LAYOUT_MAX_ENUM)
			{
				imageMemoryBarrier.oldLayout = mTexture->GetImageLayout(i);
			}
			if (imageMemoryBarrier.subresourceRange.aspectMask == VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM)
			{
				imageMemoryBarrier.subresourceRange = mTexture->GetSubresourceRange();
			}
			imageMemoryBarrier.image = mTexture->GetImage(i);
			vkCmdPipelineBarrier(
				commandBuffer,
				mSrcStageMask,
				mDstStageMask,
				0,
				0, nullptr,
				0, nullptr,
				1, &imageMemoryBarrier
			);
			mTexture->mImageLayouts[i] = imageMemoryBarrier.newLayout;
		}
	}
	else
	{
		VkImageMemoryBarrier imageMemoryBarrier = mImageMemoryBarrier;
		if (imageMemoryBarrier.oldLayout == VK_IMAGE_LAYOUT_MAX_ENUM)
		{
			imageMemoryBarrier.oldLayout = mTexture->GetImageLayout();
		}
		if (imageMemoryBarrier.subresourceRange.aspectMask == VK_IMAGE_ASPECT_FLAG_BITS_MAX_ENUM)
		{
			imageMemoryBarrier.subresourceRange = mTexture->GetSubresourceRange();
		}
		imageMemoryBarrier.image = mTexture->GetImage();
		vkCmdPipelineBarrier(
			commandBuffer,
			mSrcStageMask,
			mDstStageMask,
			0,
			0, nullptr,
			0, nullptr,
			1, &imageMemoryBarrier
		);
		mTexture->mImageLayouts[mTexture->GetFrameIndex()] = imageMemoryBarrier.newLayout;
	}
}

BufferTransfer::BufferTransfer()
{
	mBufferMemoryBarrier = {};
	mBufferMemoryBarrier.sType = VK_STRUCTURE_TYPE_BUFFER_MEMORY_BARRIER;
	mBufferMemoryBarrier.srcAccessMask = 0;
	mBufferMemoryBarrier.dstAccessMask = 0;
	mBufferMemoryBarrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	mBufferMemoryBarrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
	mBufferMemoryBarrier.offset = 0;
	mBufferMemoryBarrier.size = VK_WHOLE_SIZE;

	mSrcStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
	mDstStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
}

VkBufferMemoryBarrier BufferTransfer::GetBufferMemoryBarrier()
{
	VkBufferMemoryBarrier bufferMemoryBarrier = mBufferMemoryBarrier;
	bufferMemoryBarrier.buffer = mBuffer->GetBuffer();
	return bufferMemoryBarrier;
}

BufferTransfer& BufferTransfer::SetBuffer(RefCountPtr<RHIBuffer> buffer) noexcept
{
	mBuffer = buffer;
	return *this;
}

BufferTransfer& BufferTransfer::SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept
{
	mSrcStageMask = srcStageMask;
	return *this;
}

BufferTransfer& BufferTransfer::SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept
{
	mDstStageMask = dstStageMask;
	return *this;
}

BufferTransfer& BufferTransfer::SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept
{
	mBufferMemoryBarrier.srcAccessMask = srcAccessMask;
	return *this;
}

BufferTransfer& BufferTransfer::SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept
{
	mBufferMemoryBarrier.dstAccessMask = dstAccessMask;
	return *this;
}

BufferTransfer& BufferTransfer::SetSrcQueueFamily(VkQueueFlagBits srcQueueFlag) noexcept
{
	mBufferMemoryBarrier.srcQueueFamilyIndex = GVulkanRHI->GetQueueFamily(srcQueueFlag);
	return *this;
}

BufferTransfer& BufferTransfer::SetDstQueueFamily(VkQueueFlagBits dstQueueFlag) noexcept
{
	mBufferMemoryBarrier.dstQueueFamilyIndex = GVulkanRHI->GetQueueFamily(dstQueueFlag);
	return *this;
}

BufferTransfer& BufferTransfer::SetOffset(VkDeviceSize offset) noexcept
{
	mBufferMemoryBarrier.offset = offset;
	return *this;
}

BufferTransfer& BufferTransfer::SetSize(VkDeviceSize size) noexcept
{
	mBufferMemoryBarrier.size = size;
	return *this;
}

void BufferTransfer::Transfer(VkQueueFlagBits queueFlag /*= VK_QUEUE_GRAPHICS_BIT*/, bool isMultiple /*= false*/)
{
	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand(queueFlag);

	Transfer(commandBuffer, isMultiple);

	GRenderState->FlushSingleCommand(commandBuffer);
}

void BufferTransfer::Transfer(VkCommandBuffer commandBuffer, bool isMultiple /*= false*/)
{
	if (isMultiple && mBuffer->IsMultiple())
	{
		std::vector<VkBufferMemoryBarrier> bufferMemoryBarriers;
		bufferMemoryBarriers.resize(GRenderState->GetFrameCount());

		for (uint32 i = 0; i < GRenderState->GetFrameCount(); ++i)
		{
			bufferMemoryBarriers[i] = mBufferMemoryBarrier;
			bufferMemoryBarriers[i].buffer = mBuffer->GetBuffer(i);
		}

		vkCmdPipelineBarrier(
			commandBuffer,
			mSrcStageMask,
			mDstStageMask,
			0,
			0, nullptr,
			GRenderState->GetFrameCount(),
			bufferMemoryBarriers.data(),
			0, nullptr
		);
	}
	else
	{
		VkBufferMemoryBarrier bufferMemoryBarrier = mBufferMemoryBarrier;
		bufferMemoryBarrier.buffer = mBuffer->GetBuffer();

		vkCmdPipelineBarrier(
			commandBuffer,
			mSrcStageMask,
			mDstStageMask,
			0,
			0, nullptr,
			1, 
			&bufferMemoryBarrier,
			0, nullptr
		);
	}
}

MemoryTransfer::MemoryTransfer()
{
	mMemoryBarrier = {};
	mMemoryBarrier.sType = VK_STRUCTURE_TYPE_MEMORY_BARRIER;
	mMemoryBarrier.srcAccessMask = 0;
	mMemoryBarrier.dstAccessMask = 0;

	mSrcStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
	mDstStageMask = VK_PIPELINE_STAGE_ALL_COMMANDS_BIT;
}

VkMemoryBarrier MemoryTransfer::GetMemoryBarrier()
{
	return mMemoryBarrier;
}

MemoryTransfer& MemoryTransfer::SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept
{
	mSrcStageMask = srcStageMask;
	return *this;
}

MemoryTransfer& MemoryTransfer::SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept
{
	mDstStageMask = dstStageMask;
	return *this;
}

MemoryTransfer& MemoryTransfer::SetSrcAccessMask(VkAccessFlags srcAccessMask) noexcept
{
	mMemoryBarrier.srcAccessMask = srcAccessMask;
	return *this;
}

MemoryTransfer& MemoryTransfer::SetDstAccessMask(VkAccessFlags dstAccessMask) noexcept
{
	mMemoryBarrier.dstAccessMask = dstAccessMask;
	return *this;
}

void MemoryTransfer::Transfer(VkQueueFlagBits queueFlag /*= VK_QUEUE_GRAPHICS_BIT*/)
{
	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand(queueFlag);

	Transfer(commandBuffer);

	GRenderState->FlushSingleCommand(commandBuffer);
}

void MemoryTransfer::Transfer(VkCommandBuffer commandBuffer)
{
	vkCmdPipelineBarrier(
		commandBuffer,
		mSrcStageMask,
		mDstStageMask,
		0,
		1, &mMemoryBarrier,
		0, nullptr,
		0, nullptr
	);
}

PipelineTransfer& PipelineTransfer::SetSrcStageMask(VkPipelineStageFlags srcStageMask) noexcept
{
	mSrcStageMask = srcStageMask;
	return *this;
}

PipelineTransfer& PipelineTransfer::SetDstStageMask(VkPipelineStageFlags dstStageMask) noexcept
{
	mDstStageMask = dstStageMask;
	return *this;
}

PipelineTransfer& PipelineTransfer::AddMemory(MemoryTransfer& memoryTransfer)
{
	mMemoryBarriers.push_back(memoryTransfer.GetMemoryBarrier());
	return *this;
}

PipelineTransfer& PipelineTransfer::AddBuffer(BufferTransfer& bufferTransfer)
{
	mBufferMemoryBarriers.push_back(bufferTransfer.GetBufferMemoryBarrier());
	return *this;
}

PipelineTransfer& PipelineTransfer::AddTexture(TextureTransfer& textureTransfer)
{
	mImageMemoryBarriers.push_back(textureTransfer.GetImageMemoryBarrier());
	return *this;
}

void PipelineTransfer::Transfer(VkQueueFlagBits queueFlag /*= VK_QUEUE_GRAPHICS_BIT*/)
{
	VkCommandBuffer commandBuffer = GRenderState->BeginSingleCommand(queueFlag);

	Transfer(commandBuffer);

	GRenderState->FlushSingleCommand(commandBuffer);
}

void PipelineTransfer::Transfer(VkCommandBuffer commandBuffer)
{

	vkCmdPipelineBarrier(
		commandBuffer,
		mSrcStageMask,
		mDstStageMask,
		0,
		mMemoryBarriers.size(), mMemoryBarriers.data(),
		mBufferMemoryBarriers.size(), mBufferMemoryBarriers.data(),
		mImageMemoryBarriers.size(), mImageMemoryBarriers.data()
	);
}
